/********************************************************
 * Based on code from http://system.data.sqlite.org
 *
 *  Author 杨岑.
 *  Copyright (c) 2015 杨岑. All rights reserved.
 *
 * The original code is for Windows platform only.
 * I replaced the crypto part to use MacOS system functions.
 ********************************************************/


#ifndef SQLITE_OMIT_DISKIO
#ifdef SQLITE_HAS_CODEC

#include <CommonCrypto/CommonDigest.h>
#include <CommonCrypto/CommonCryptor.h>
#include "sqlite3.c"


#define CTR_CIPHER_ALGO       kCCAlgorithmAES       /* Use AES */
#define CTR_CIPHER_BLOCK_SIZE kCCBlockSizeAES128    /* AES only supports 128-bit block */
#define CTR_CIPHER_KEY_SIZE   kCCKeySizeAES256      /* Use wide 32-byte key size; non-used bits are zero'ed out */
#define CTR_CIPHER_MODE       kCCOptionECBMode      /* CTR encryption mode utilizes ECB mode */
#define PASSWORD_HASH_REPEAT  1000000               /* Repeat how many times to derive key from password */

static const char* CIPHER_ERROR = "Unable to create cryptor!";

/*
 * This structure contains critical data for encryption/decryption of pager data.
 */
typedef struct CipherContext
{
    CCCryptorRef hReadKey;  /* Key used to read from the database and write to the journal */
    CCCryptorRef hWriteKey; /* Key used to write to the database */
    void*   pbk;            /* Password hash */
    int     pageSize;       /* Size of pages */
    void*	pageBuffer;     /* Buffer for page encryption/decryption */
} CipherContext;



/*
 * This function implements the counter mode (CTR) encryption.
 * In order to be independent of future sqlite changes, it uses offset
 * of the data as the base of counter calculation (offset/block size).
 * In CTR mode, encryption and decryption are same process.
 */
CCCryptorStatus cryptoCounterModeCrypt(CCCryptorRef cipher, void* data, size_t dataLength, int64_t counterHint)
{
    /* Nounce size must be same as block size */
    static char nounce[] = {0x5c,0x3e,0x43,0x98,0x8e,0xe9,0xb6,0x9b,0xf8,0x43,0x56,0xcd,0x3c,0xab,0x50,0x89};
    assert(sizeof(nounce) == CTR_CIPHER_BLOCK_SIZE);
    
    char counterBuffer[CTR_CIPHER_BLOCK_SIZE];
    char keystreamBuffer[CTR_CIPHER_BLOCK_SIZE];
    int64_t counter = (counterHint + 1) / CTR_CIPHER_BLOCK_SIZE;
    size_t numBlocks = (dataLength + CTR_CIPHER_BLOCK_SIZE - 1) / CTR_CIPHER_BLOCK_SIZE;
    int bytesProcessed = 0;
    
    for (int i = 0; i < numBlocks; i++)
    {
        memcpy(counterBuffer, nounce, sizeof(nounce));
        *(int64_t*)counterBuffer ^= counter++;

        size_t moved = 0;
        CCCryptorReset(cipher, 0);
        CCCryptorStatus error = CCCryptorUpdate(cipher, counterBuffer, CTR_CIPHER_BLOCK_SIZE, keystreamBuffer, CTR_CIPHER_BLOCK_SIZE, &moved);
        assert(error == kCCSuccess);
        assert(moved == CTR_CIPHER_BLOCK_SIZE);
        
        if (error == kCCSuccess)
        {
            for (int j = 0; j < CTR_CIPHER_BLOCK_SIZE && bytesProcessed < dataLength; j++)
            {
                ((char*)data)[bytesProcessed++] ^= keystreamBuffer[j];
            }
        }
        else return error;
    }

    return kCCSuccess;
}


/*
 * Creates or updates a cryptographic context for a pager.
 */
CipherContext* cryptoCreateCipherContext(CCCryptorRef hKey, int pageSize, CipherContext* existingContext)
{
	CipherContext* context;

    if (existingContext)
    {
        context = existingContext;
        sqlite3_free(context->pageBuffer);
    }
    else
    {
        context = (CipherContext*)sqlite3_malloc(sizeof(CipherContext));
        if (!context) return NULL;

        memset(context, 0, sizeof(CipherContext));
        context->hReadKey = hKey;
        context->hWriteKey = hKey;
    }

    context->pageSize = pageSize;
    context->pageBuffer = sqlite3_malloc(pageSize);

    return context;
}

/*
 * Derives key hash from password. Key size is CTR_CIPHER_KEY_SIZE.
 *
 * Don't forget to free the returned pointer.
 */
void* cryptoDerivePasswordHash(const void* pKey, int nKeyLen)
{
    static const char salt[] = {
        0xa2,0x24,0x89,0x48,0xdf,0x38,0x60,0x2c,0xc3,0x08,0xd6,0xd1,0x23,0x89,0x79,0x36,
        0x14,0xa6,0xcd,0x32,0xa1,0x4e,0xb4,0x40,0x94,0x45,0x79,0xfe,0xac,0x9c,0x45,0xa8,
        0xd4,0x37,0x40,0xaf,0x71,0xe8,0x13,0x61,0xc7,0x68,0x55,0x71,0x51,0x75,0x4b,0xbf,
        0xc4,0xd0,0xeb,0x84,0xb4,0xff,0xbf,0xca,0x8e,0x48,0x94,0xf8,0xd5,0x86,0x44,0x0e
    };
    
    int passwordDataSize = sizeof(salt) + nKeyLen;
    char* passwordData = sqlite3_malloc(passwordDataSize);
    if (!passwordData) return 0;
    
    const int keySize = CC_SHA256_DIGEST_LENGTH;    /* Same as key size */
    assert(keySize == CTR_CIPHER_KEY_SIZE);
    unsigned char* keyHash = sqlite3_malloc(keySize);
    if (!keyHash) return 0;
    
    memcpy(passwordData, pKey, nKeyLen);
    memcpy(passwordData + nKeyLen, salt, sizeof(salt));
    
    CC_SHA256(passwordData, (CC_LONG)passwordDataSize, keyHash);
    sqlite3_free(passwordData);
    
    int i;
    for(i = 0; i < PASSWORD_HASH_REPEAT; i++)
    {
        CC_SHA256(keyHash, keySize, keyHash);
    }
    return keyHash;
}


CCCryptorRef cryptoCreateCipher(const void *key, int keySize)
{
    MUTEX_LOGIC( sqlite3_mutex *pMaster = sqlite3_mutex_alloc(SQLITE_MUTEX_STATIC_MASTER); )
    sqlite3_mutex_enter(pMaster);

    CCCryptorRef cipher = 0;
    CCCryptorStatus error = CCCryptorCreate(kCCEncrypt, CTR_CIPHER_ALGO, CTR_CIPHER_MODE, key, keySize, 0, &cipher);
    sqlite3_mutex_leave(pMaster);
    assert(error == kCCSuccess);
    
    return error == kCCSuccess ? cipher : 0;
}


/**********************************************************************/


/* Needed for re-keying */
static void * sqlite3pager_get_codecarg(Pager *pPager)
{
    return (pPager->xCodec) ? pPager->pCodec: NULL;
}

void sqlite3_activate_see(const char *info)
{
}


/* Destroy a cryptographic context and any buffers and keys allocated therein */
static void sqlite3CodecFree(void* pv)
{
    CipherContext* cipherContext = (CipherContext*)pv;
    /* Destroy the read key if there is one */
    if (cipherContext->hReadKey)
    {
        CCCryptorRelease(cipherContext->hReadKey);
    }

    /* If there's a writekey and its not equal to the readkey, destroy it */
    if (cipherContext->hWriteKey && cipherContext->hWriteKey != cipherContext->hReadKey)
    {
        CCCryptorRelease(cipherContext->hWriteKey);
    }

    sqlite3_free(cipherContext->pbk);
    sqlite3_free(cipherContext->pageBuffer);
    sqlite3_free(cipherContext);
}


void sqlite3CodecSizeChange(void *pArg, int pageSize, int reservedSize)
{
    CipherContext* cipherContext = (CipherContext*)pArg;

    if (cipherContext->pageSize != pageSize)
    {
        cryptoCreateCipherContext(cipherContext->hReadKey, pageSize, cipherContext);
    }
}


/* Encrypt/Decrypt functionality, called by pager.c */
void * sqlite3Codec(void *pArg, void *data, Pgno nPageNum, int nMode)
{
    CipherContext* cipherContext = (CipherContext*)pArg;

    if (!cipherContext) return data;

    CCCryptorRef cryptor = 0;
    void* pData = 0;
    switch(nMode)
    {
        case 0: /* Undo a "case 7" journal file encryption */
        case 2: /* Reload a page */
        case 3: /* Load a page */
        	cryptor = cipherContext->hReadKey;
        	pData = data;
            break;
        case 6: /* Encrypt a page for the main database file */
        	cryptor = cipherContext->hWriteKey;
        	pData = cipherContext->pageBuffer;
        	memcpy(pData, data, cipherContext->pageSize);
            break;
        case 7: /* Encrypt a page for the journal file */
            /* Under normal circumstances, the readkey is the same as the writekey.  However,
             when the database is being rekeyed, the readkey is not the same as the writekey.
             The rollback journal must be written using the original key for the
             database file because it is, by nature, a rollback journal.
             Therefore, for case 7, when the rollback is being written, always encrypt using
             the database's readkey, which is guaranteed to be the same key that was used to
             read the original data.
             */
        	cryptor = cipherContext->hReadKey;
        	pData = cipherContext->pageBuffer;
        	memcpy(pData, data, cipherContext->pageSize);
        	break;
    }
    
    if (cryptor)
    {
    	cryptoCounterModeCrypt(cryptor, pData, cipherContext->pageSize, cipherContext->pageSize * (nPageNum - 1));
        return pData;
    }
    else
    	return data;
}


/* Called by sqlite and sqlite3_key_interop to attach a key to a database. */
int sqlite3CodecAttach(sqlite3 *db, int nDb, const void *pKey, int nKeyLen)
{
    int rc = SQLITE_ERROR;
    CCCryptorRef hKey = 0;
    char* pbk = 0;

    /* No key specified, could mean either use the main db's encryption or no encryption */
    if (!pKey || !nKeyLen)
    {
        if (!nDb)
        {
            return SQLITE_OK; /* Main database, no key specified so not encrypted */
        }
        else /* Attached database, use the main database's key */
        {
            /* Get the encryption block for the main database and attempt to duplicate the key
             ** for use by the attached database
             */
            Pager *p = sqlite3BtreePager(db->aDb[0].pBt);
            CipherContext* cipherContext = (CipherContext*)sqlite3pager_get_codecarg(p);

            if (!cipherContext) return SQLITE_OK; /* Main database is not encrypted so neither will be any attached database */
            if (!cipherContext->hReadKey) return SQLITE_OK; /* Not encrypted */

            pbk = sqlite3_malloc(CTR_CIPHER_KEY_SIZE);
            if (!pbk) return SQLITE_NOMEM;
            memcpy(pbk, cipherContext->pbk, CTR_CIPHER_KEY_SIZE);

            // Duplicate the key by create a new cipher with the same password book
            hKey = cryptoCreateCipher(pbk, CTR_CIPHER_KEY_SIZE);
            if (!hKey)
            {
            	sqlite3ErrorWithMsg(db, rc, CIPHER_ERROR);
            	return rc; /* Unable to duplicate the key */
            }
        }
    }
    else /* User-supplied passphrase, so create a cryptographic key out of it */
    {
        pbk = cryptoDerivePasswordHash(pKey, nKeyLen);
        if (!pbk) return SQLITE_NOMEM;
        
        hKey = cryptoCreateCipher(pbk, CTR_CIPHER_KEY_SIZE);
        if (!hKey)
        {
        	sqlite3_free(pbk);
        	sqlite3ErrorWithMsg(db, rc, CIPHER_ERROR);
            return rc;
        }
    }

    /* Create a new encryption block and assign the codec to the new attached database */
    if (hKey)
    {
        Pager *p = sqlite3BtreePager(db->aDb[nDb].pBt);
        CipherContext* cipherContext = cryptoCreateCipherContext(hKey, p->pageSize, NULL);
        if (!cipherContext)
        {
        	sqlite3_free(pbk);
        	CCCryptorRelease(hKey);
        	return SQLITE_NOMEM;
        }
        
		cipherContext->pbk = pbk;
		sqlite3PagerSetCodec(p, sqlite3Codec, sqlite3CodecSizeChange, sqlite3CodecFree, cipherContext);

		rc = SQLITE_OK;
    }
    return rc;
}


/* Once a password has been supplied and a key created, we don't keep the
** original password for security purposes.  Therefore return NULL.
*/
void sqlite3CodecGetKey(sqlite3 *db, int nDb, void **ppKey, int *pnKeyLen)
{
    Btree *pbt = db->aDb[0].pBt;
    Pager *p = sqlite3BtreePager(pbt);
    CipherContext* cipherContext = (CipherContext*)sqlite3pager_get_codecarg(p);

    if (ppKey) *ppKey = 0;
    if (pnKeyLen && cipherContext) *pnKeyLen = 0;
}


/* Changes the encryption key for an existing database. */
SQLITE_API int sqlite3_rekey_v2(sqlite3 *db, const char* zDbName, const void* pKey, int nKeySize)
{
    Btree *pbt = db->aDb[0].pBt;
    Pager *p = sqlite3BtreePager(pbt);
    CipherContext* cipherContext = (CipherContext*)sqlite3pager_get_codecarg(p);
    CCCryptorRef hKey;
    void* pbk;
    int rc = SQLITE_ERROR;

    pbk = cryptoDerivePasswordHash(pKey, nKeySize);
    if (!pbk) return SQLITE_NOMEM;
    
	hKey = cryptoCreateCipher(pbk, CTR_CIPHER_KEY_SIZE);
    if (hKey == 0)
    {
        sqlite3_free(pbk);
        sqlite3ErrorWithMsg(db, rc, CIPHER_ERROR);
        return rc;
    }

    if (!cipherContext && !hKey) return SQLITE_OK; /* Wasn't encrypted to begin with */

    /* To rekey a database, we change the writekey for the pager.
     * The readkey remains the same.
     */
    if (!cipherContext) /* Encrypt an unencrypted database */
    {
        cipherContext = cryptoCreateCipherContext(hKey, p->pageSize, NULL);
        if (!cipherContext)
            return SQLITE_NOMEM;

        cipherContext->hReadKey = 0; /* Original database is not encrypted */
        sqlite3PagerSetCodec(sqlite3BtreePager(pbt), sqlite3Codec, sqlite3CodecSizeChange, sqlite3CodecFree, cipherContext);
    }
    else /* Change the writekey for an already-encrypted database */
    {
        cipherContext->hWriteKey = hKey;
    }

    sqlite3_mutex_enter(db->mutex);

    /* Start a transaction */
    rc = sqlite3BtreeBeginTrans(pbt, 1, NULL);

    if (!rc)
    {
        /* Rewrite all the pages in the database using the new encryption key */
        Pgno nPage;
        Pgno nSkip = PAGER_MJ_PGNO(p);
        DbPage *pPage;
        Pgno n;
        int count;

        sqlite3PagerPagecount(p, &count);
        nPage = (Pgno)count;

        for(n = 1; n <= nPage; n ++)
        {
            if (n == nSkip) continue;
            rc = sqlite3PagerGet(p, n, &pPage, 0);
            if(!rc)
            {
                rc = sqlite3PagerWrite(pPage);
                sqlite3PagerUnref(pPage);
            }
        }
    }

    /* If we succeeded, try and commit the transaction */
    if (!rc)
    {
        rc = sqlite3BtreeCommit(pbt);
    }

    /* If we failed, rollback */
    if (rc)
    {
        sqlite3BtreeRollback(pbt, SQLITE_OK, FALSE);
    }

    /* If we succeeded, destroy any previous read key this database used
     ** and make the readkey equal to the writekey
     */
    if (!rc)
    {
        if (cipherContext->hReadKey)
        {
            CCCryptorRelease(cipherContext->hReadKey);
        }
        sqlite3_free(cipherContext->pbk);
        cipherContext->hReadKey = cipherContext->hWriteKey;
    }
    /* We failed.  Destroy the new writekey (if there was one) and revert it back to
     ** the original readkey
     */
    else
    {
    	sqlite3_free(pbk);
        if (cipherContext->hWriteKey)
        {
            CCCryptorRelease(cipherContext->hWriteKey);
        }
        cipherContext->hWriteKey = cipherContext->hReadKey;
    }

    /* If the readkey and writekey are both empty, there's no need for a codec on this
     ** pager anymore.  Destroy the crypt block and remove the codec from the pager.
     */
    if (!cipherContext->hReadKey && !cipherContext->hWriteKey)
    {
        sqlite3PagerSetCodec(p, NULL, NULL, NULL, NULL);
    }
    
    sqlite3_mutex_leave(db->mutex);
    
    return rc;
}

SQLITE_API int sqlite3_rekey(sqlite3 *db, const void* pKey, int nKeySize)
{
    return sqlite3_rekey_v2(db, 0, pKey, nKeySize);
}

SQLITE_API int sqlite3_key_v2(sqlite3 *db, const char* zDbName, const void* pKey, int nKeySize)
{
    return sqlite3CodecAttach(db, 0, pKey, nKeySize);
}

/* We do not attach this key to the temp store, only the main database. */
SQLITE_API int sqlite3_key(sqlite3 *db, const void* pKey, int nKeySize)
{
    return sqlite3_key_v2(db, 0, pKey, nKeySize);
}

#endif /* SQLITE_HAS_CODEC */
#endif /* SQLITE_OMIT_DISKIO */


